% w is a sequence of words making up the corpus
% z is a sequence of labels for each word
% alpha is a vector length K

% M = number of documents (400)
% K = number of topics (3)
% V = number of terms t in vocabulary (6000?)

%k and z are somehow used interchangeably for a topic - maybe k is a number
%z is an actual label?

% Initialization:
%n_m^(k) : document-topic count: number of times topic k appears with any word in document m
%n_k^(t) : topic-term count: number of times word t has topic k(?)
%n_m : document-topic sum - sum of all n_m^(k)?
%n_k : topic-term sum - sum of all n_k^(t)?
load('classicer527.mat');
K = 4;

n_mk = zeros(size(classic400,1),K); % probably shouldn't hard code the 3 (number of labels)
maxdoc = 0;
for m=1:size(classic400,1)
    docsize = sum(classic400(m,classic400(m,:)>0));
    if(maxdoc < docsize)
        maxdoc = docsize;
    end
end
n_kt = zeros(K,size(classic400,2));
z = zeros(size(classic400,1),maxdoc);
for m=1:size(z,1) %for all documents m in document pool M
    n = 1;
    for t=find(classic400(m,:)>0) %for all words n in document m
        %Mult(1/K) - sample assuming all topics are equally probable?
        for count=1:classic400(m,t)
            k = randi([1,K],1);
            z(m,n) = k;
            n = n + 1;
            n_mk(m,k) = n_mk(m,k) + 1;
            n_kt(k,t) = n_kt(k,t) + 1;
        end
        
    end
end
alpha = ones(1,K);
beta = ones(1,size(classic400,2));

phi_init = read_phi(n_kt,beta);
theta_init = read_theta(n_mk,alpha);
%Gibbs sampling over burn-in period and sampling period
for counter=1:20
    counter
    for m=1:size(z,1) %for all documents m in document pool M
        n = 1;
        %[counter,m]
        for t=find(classic400(m,:)>0) %for all words n in document m            
            for count=1:classic400(m,t)
                n_mk(m,z(m,n)) = n_mk(m,z(m,n)) - 1;
                n_kt(z(m,n),t) = n_kt(z(m,n),t) - 1;

                k = gibbs_sample(m,t,n_mk,n_kt,alpha,beta); % i = m,n

                z(m,n) = k;
                n_mk(m,k) = n_mk(m,k) + 1;
                n_kt(k,t) = n_kt(k,t) + 1;
                n = n + 1;
            end
        end
    end
end
%read out params
phi = read_phi(n_kt,beta);
theta = read_theta(n_mk,alpha);

topic1 = zeros(1,size(phi,2));
topic2 = zeros(1,size(phi,2));
topic3 = zeros(1,size(phi,2));
topic4 = zeros(1,size(phi,2));
for i=1:size(phi,2)
    [c,idx] = max(phi(:,i));
    if idx == 1
        topic1(i) = c;
    elseif idx == 2
        topic2(i) = c;
    elseif idx == 3
        topic3(i) = c;
    elseif idx == 4
        topic4(i) = c;
    end
end

[~,ix1] = sort(topic1,'descend');   
[~,ix2] = sort(topic2,'descend');
[~,ix3] = sort(topic3,'descend');
[~,ix4] = sort(topic4,'descend');

for i=1:15
    words1{i} = classicwordlist{ix1(i)};
end

for i=1:15
    words2{i} = classicwordlist{ix2(i)};
end

for i=1:15
    words3{i} = classicwordlist{ix3(i)};
end

for i=1:15
    words4{i} = classicwordlist{ix4(i)};
end

X = theta(1:111,1);
X_ = theta(1:111,2);
Y = theta(1:111,3);
Y_ = theta(1:111,4);
plot(X-X_,Y-Y_,'.');
hold on;
X = theta(112:248,1);
X_ = theta(112:248,2);
Y = theta(112:248,3);
Y_ = theta(112:248,4);
p1 = plot(X-X_,Y-Y_,'.');
set(p1,'Color','red');
hold on;
X = theta(249:389,1);
X_ = theta(249:389,2);
Y = theta(249:389,3);
Y_ = theta(249:389,4);
p2 = plot(X-X_,Y-Y_,'.');
set(p2,'Color','green');
hold on;
X = theta(390:527,1);
X_ = theta(390:527,2);
Y = theta(390:527,3);
Y_ = theta(390:527,4);
p2 = plot(X-X_,Y-Y_,'.');
set(p2,'Color','yellow');
grid on;
axis square;